import torch
import torch.nn as nn
import torch.optim as optim

from data_manager.loaders.data_loader import get_loader
from models_manger import GoogleNet
from trainer_manger.trainer_manger import TrainerManger

torch.manual_seed(0)
"""
控制随机初始化
    影响 权重初始化（如 nn.Linear、nn.Conv2d 的初始参数）。
    影响 数据打乱（如 DataLoader(shuffle=True)）。
    影响 Dropout、随机增强（如 transforms.RandomCrop）等涉及随机性的操作。"""


def test():
    model = GoogleNet(num_class=102)
    print(model)
    criterion = nn.CrossEntropyLoss()  # 交叉熵损失
    optimizer = optim.SGD(model.parameters(), lr=0.002, momentum=0.9)

    train_loader, test_loader = get_loader("flowers102", batch_size=64)

    epoch = 200
    # Adam优化器
    trainer_manger = TrainerManger(model, optimizer, criterion, train_loader, test_loader, epoch)
    trainer_manger.train()


if __name__ == '__main__':
    test()
